from core.function import Function
from core.cuda import cuda_module

# 对输入数据进行裁剪
# 这种操作通常用于梯度裁剪（Gradient Clipping）等场景，其中梯度值可能过大，可能导致梯度爆炸的问题。
# 通过对梯度进行裁剪，可以防止梯度过大，有助于稳定训练过程

class Clip(Function):
    def __init__(self, x_min, x_max):
        """
        初始化 Clip 函数。

        参数:
            x_min (float): 进行剪裁的最小值。
            x_max (float): 进行剪裁的最大值。
        """
        self.x_min = x_min
        self.x_max = x_max

    def forward(self, x):
        """
        Clip 函数的前向传播。

        参数:
            x (ndarray): 待剪裁的输入数组。

        返回:
            ndarray: 剪裁后的数组。
        """
        xp = cuda_module
        y = xp.clip(x, self.x_min, self.x_max)
        return y

    def backward(self, gy):
        """
        Clip 函数的反向传播。

        参数:
            gy (ndarray): 输出的梯度。

        返回:
            ndarray: 相对于输入的梯度。
        """
        x, = self.inputs
        mask = (x.data >= self.x_min) * (x.data <= self.x_max)
        gx = gy * mask
        return gx

    def clip(self, x):
        """
        剪裁操作的接口。

        参数:
            x (ndarray): 待剪裁的输入数组。

        返回:
            ndarray: 剪裁后的数组。
        """
        return self(x)

